import torch
from torch.utils.data import DataLoader
import torchvision
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torchvision.datasets import FashionMNIST, CIFAR10


def load_fashion_mnist(batch_size, reshape=None, root='Data'):

    if reshape:
        transform = transforms.Compose([
            transforms.Resize(reshape),
            transforms.ToTensor(),
        ])
    else:
        transform = transforms.Compose([
            transforms.ToTensor(),
        ])
    
    train_data = FashionMNIST(root=root, train=True, download=True, transform=transform)
    test_data = FashionMNIST(root=root, train=False, download=True, transform=transform)
    train_data_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=4)
    test_data_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=4)

    return train_data_loader, test_data_loader, 10, 1


def load_cifar10(batch_size, reshape=None, root='Data/Cifar10'):

    if reshape:
        transform = transforms.Compose([
            transforms.Resize(reshape),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])
    else:
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
        ])
    
    train_dataset = CIFAR10(root='Data/Cifar10', train=True, download=True, transform=transform)
    test_dataset = CIFAR10(root='Data/Cifar10', train=False, download=True, transform=transform)
    train_data_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    test_data_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    return train_data_loader, test_data_loader, 10, 3


def load_awa2(batch_size, resize=None, root='Data/AwA2/'):
    train_path = root + "train"
    test_path = root + "test"
    train_transform= transforms.Compose([
        transforms.RandomCrop(224, 8),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    train_dataset = ImageFolder(train_path, transform=train_transform)
    test_dataset = ImageFolder(test_path, transform=test_transform)
    train_data_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    test_data_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
    return train_data_loader, test_data_loader, 50, 3


if __name__ == '__main__':
    train_data_loader, test_data_loader = load_awa2(128)
    for i, (images, labels) in enumerate(train_data_loader):
        print(images.shape, labels.shape)
        exit(0)
    